#!/usr/bin/python3
# -*- coding: utf-8 -*-

#publish images
import rospy
from sensor_msgs.msg import Image
import gxipy as gx
from cv_bridge import CvBridge,CvBridgeError
import cv2
import sys

import time
import os
import numpy as np
import traceback

np.set_printoptions(formatter={'float': '{: 0.4f}'.format})

import message_filters

from geometry_msgs.msg import Transform,Pose,Point32

from bs_imu import bs_imu
from bs_img import bs_img_real
from bs_img import bs_img_base
from bs_img.bs_cfg_real_5m import *
from bs_img.bs_lie_solve import *
from bs_img.serial_pack import *





class BS_Solver:
    def __init__(self):
        self.last_plane_pts = np.array([])   #保存上一帧点用于配对
        self.is_first_frame = True                  #判断是否为第一帧

        self.invalid_frame_num_plane = 0      #compute continus invalid frame
        self.invalid_frame_num_drogue = 0
        self.invalid_frame_max = 5     #if invaild frame < invalid_frame_max, output results of the last frame

        #Kalman Filter
        self.Q = np.eye(3)
        self.R = np.eye(3)
        self.x_plane = np.array([np.nan, np.nan, np.nan])
        self.x_drogue = np.array([np.nan, np.nan, np.nan])
        self.P_plane = np.eye(3)
        self.P_drogue = np.eye(3)
        
        self.lf_img_pub = rospy.Publisher('image_dh/lf', Image, queue_size=1)
        self.lf_plane_pub = rospy.Publisher("/bs_debuger/result/lf_plane", Transform , queue_size=1)
        self.lf_drogue_pub = rospy.Publisher("/bs_debuger/result/lf_drogue", Transform, queue_size=1)
        
        self.bridge = CvBridge()

        
        device_manager = gx.DeviceManager()
        dev_num, dev_info_list = device_manager.update_device_list()
        if dev_num == 0:
            sys.exit(1)
        '''
        每个相机的sn都是独特的，在相机下面写着，这样就可以利用这个分开长短焦相机
        sn作为launch文件中外传入的参数
        '''

        param_list = rospy.get_param_names()
        for param in param_list:  #反证只有一个键，虽说跟node有关
            if '/sn' in param:
                sn = rospy.get_param(param) 
            if '/ExposureTime' in param:
                ExposureTime = rospy.get_param(param)
        self.cam = device_manager.open_device_by_sn(sn)
        #set parameters
        self.cam.TriggerMode.set(gx.GxSwitchEntry.OFF)#set continuous acquisition
        self.cam.ExposureTime.set(ExposureTime)  #set exposure time
        self.cam.stream_on()
        
        self.solve()
        
        if rospy.is_shutdown():
            self.cam.stream_off()
            self.cam.close_device()



    def lf_img_cb(self,img_gray):
        _, img_bin = cv2.threshold(img_gray, 70, 0xff, cv2.THRESH_BINARY)
        # 如果相机倒置需要旋转一下
        # lf_img_gray = cv2.rotate(lf_img_gray, cv2.ROTATE_180)
        plane_t_vec = [np.nan]*3
        drogue_t_vec = [np.nan]*3
        drogue_pts = []
        
        gt_pts,aera_list = bs_img_real.img_to_pts(img_bin)
        gt_pts = np.array(gt_pts)
        cnt_plane_pts = [ ]
        
        len_gt_pts = len(gt_pts)
        if len_gt_pts >= 18:    #噪点太多（可能逆光），为了防止其解算时间太长耽误后面的进度，这里直接跳过
            return


        if self.is_first_frame:   #第一帧图像
            drogue_result_dict = bs_img_real.solve_drogue(gt_pts, camK_lf, taper_radius=r_drogue)        
            if drogue_result_dict["img_valid"]:
                drogue_t_vec = drogue_result_dict["t_vec"]
                drogue_pts = drogue_result_dict["drogue_px_calcL"]
            else:
                drogue_pts = []
            drogue_pts = np.array(drogue_pts)
            if np.size(drogue_pts) >0:
                ind = ~np.isin(gt_pts[:, 0], drogue_pts[:, 0]) | ~np.isin(gt_pts[:, 1], drogue_pts[:, 1])
                plane_pts = np.array(gt_pts)[ind]
                ind2 = ind == False
                plane_aeras = np.array(aera_list)[ind2] 
            else:
                plane_pts = gt_pts
                plane_aeras = np.array(aera_list)
            if len(plane_pts)>=5:
                plane_result_dict = bs_img_base.solve_plane_pt_new(plane_pts,plane_real_ptL,cam_K=camK_lf,distort_coeffs=distortion_coeffs_lf,plane_area_ptL=plane_aeras)
                if plane_result_dict["img_valid"]:
                    self.is_first_frame = False 
                    cnt_plane_pts = plane_result_dict["plane_px_calcL"]
                    plane_t_vec = plane_result_dict["t_vec"].flatten()
                    plane_quat = bs_imu.rot_vec_to_quat(plane_result_dict["r_vec"].flatten())

        else:   #除第一帧以外的图像
            match_pt_idxs = find_matched_pts(self.last_plane_pts,-np.mean(gt_pts,axis=0))
            flag_empty = True   #指示匹配是否为空
            plane_aeras = []
            for match_pt_idx in match_pt_idxs:
                if match_pt_idx is not None:
                    cnt_plane_pts.append(gt_pts[match_pt_idx])
                    plane_aeras.append(gt_pts[match_pt_idx])
                    flag_empty = False
            if flag_empty:
                self.is_first_frame = True   #匹配为空重新开始，视为第一帧
            
            # 至少有5个点则解算飞机特征点
            if len(cnt_plane_pts) >= 5:    
                plane_result_dict = bs_img_base.solve_plane_pt_new(cnt_plane_pts,plane_real_ptL,cam_K=camK_lf,distort_coeffs=distortion_coeffs_lf,plane_area_ptL=plane_aeras)
                if plane_result_dict["img_valid"]:
                    self.is_first_frame = False 
                    cnt_plane_pts = plane_result_dict["plane_px_calcL"]
                    plane_t_vec = plane_result_dict["t_vec"].flatten()
                else:
                    cnt_plane_pts = []
            else:
                cnt_plane_pts = []
            cnt_plane_pts = np.array(cnt_plane_pts)
                
            if np.size(cnt_plane_pts) >0:
                ind = ~np.isin(gt_pts[:, 0], cnt_plane_pts[:, 0]) | ~np.isin(gt_pts[:, 1], cnt_plane_pts[:, 1])
                drogue_pts = np.array(gt_pts)[ind]
            else:
                drogue_pts = gt_pts
            
            drogue_result_dict = bs_img_real.solve_drogue(drogue_pts, camK_lf, taper_radius=r_drogue)        
            if drogue_result_dict["img_valid"]:
                drogue_t_vec = drogue_result_dict["t_vec"]

        
        self.last_plane_pts = cnt_plane_pts - np.mean(cnt_plane_pts,axis=0)


        plane_is_visiable = True if not(np.all(np.isnan(plane_t_vec))) else False
        drogue_is_visiable = True if not(np.all(np.isnan(drogue_t_vec))) else False
        
        # #compensate
        # if plane_is_visiable:
        #     plane_t_vec[0] = plane_t_vec[0]/0.81*0.4235+0.0086*plane_t_vec[2]-2.9211+2.8
        #     plane_t_vec[1] = plane_t_vec[1]-0.0266*0.93*plane_t_vec[2]+3.4377-0.35-2.5
        #     plane_t_vec[2] = plane_t_vec[2]*0.97+3.68
        # if drogue_is_visiable:
        #     drogue_t_vec[0] = drogue_t_vec[0]*0.9336+0.23
        #     drogue_t_vec[1] = drogue_t_vec[1]/1.3
        #     drogue_t_vec[2] = drogue_t_vec[2]*0.9034/(2.58107+0.244865)*(1.01903+0.244865)-0.1386+0.0766

        # # KF for plane
        # if not(plane_is_visiable):
        #     self.invalid_frame_num_plane = self.invalid_frame_num_plane + 1
        #     if self.invalid_frame_num_plane <= self.invalid_frame_max:
        #         plane_res = self.x_plane
        #         if not(np.all(np.isnan(plane_res))):
        #             plane_res_frame = serial_frame("drogue", True, round(plane_res[2]*100),round(plane_res[0]*100),round(plane_res[1]*100))  
        #         else:
        #             plane_res_frame = serial_frame("drogue",False,0,0,0)
        #     else:
        #         self.x_plane = np.array([np.nan, np.nan, np.nan])
        #         plane_res_frame = serial_frame("drogue",False,0,0,0)
        
        # else:
        #     self.invalid_frame_num_plane = 0

        #     if np.all(np.isnan(self.x_plane)):   #the first frame
        #         self.x_plane = plane_t_vec
        #     else:
        #         z_measure_plane = np.array(plane_t_vec)
        #         self.P_plane = self.P_plane + self.Q
        #         K = np.matmul(self.P_plane, np.linalg.pinv(self.P_plane + self.R))
        #         self.x_plane = self.x_plane + np.matmul(K, z_measure_plane - self.x_plane)
        #         self.P_plane = np.matmul(np.eye(3) - K, self.P_plane)
        #     plane_res_frame = serial_frame("drogue",True,round(self.x_plane[2]*100),round(self.x_plane[0]*100),round(self.x_plane[1]*100))  

        # uart.write(plane_res_frame)
        # time.sleep(0.0005)

        # # KF for drogue
        # if not(drogue_is_visiable):
        #     self.invalid_frame_num_drogue = self.invalid_frame_num_drogue + 1
        #     if self.invalid_frame_num_drogue <= self.invalid_frame_max:
        #         drogue_res = self.x_drogue
        #         if not(np.all(np.isnan(drogue_res))):
        #             drogue_res_frame = serial_frame("drogue", True, round(drogue_res[2]*100),round(drogue_res[0]*100),round(drogue_res[1]*100))  
        #         else:
        #             drogue_res_frame = serial_frame("drogue",False,0,0,0)
        #     else:
        #         self.x_drogue = np.array([np.nan, np.nan, np.nan])
        #         drogue_res_frame = serial_frame("drogue",False,0,0,0)
        
        # else:
        #     self.invalid_frame_num_drogue = 0

        #     if np.all(np.isnan(self.x_drogue)):   #the first frame
        #         self.x_drogue = drogue_t_vec
        #     else:
                # z_measure_drogue = np.array(drogue_t_vec)
                # self.P_drogue = self.P_drogue + self.Q
                # K = np.matmul(self.P_drogue, np.linalg.pinv(self.P_drogue + self.R))
                # self.x_drogue = self.x_drogue + np.matmul(K, z_measure_drogue - self.x_drogue)
                # self.P_drogue = np.matmul(np.eye(3) - K, self.P_drogue)
        #     drogue_res_frame = serial_frame("drogue",True,round(self.x_drogue[2]*100),round(self.x_drogue[0]*100),round(self.x_drogue[1]*100))  

        # uart.write(drogue_res_frame)
        # time.sleep(0.0005)

        plane_msg = Transform()
        plane_msg.translation.x = plane_t_vec[0]
        plane_msg.translation.y = plane_t_vec[1]
        plane_msg.translation.z = plane_t_vec[2]
        self.lf_plane_pub.publish(plane_msg)


        drogue_msg = Transform()
        drogue_msg.translation.x = drogue_t_vec[0]
        drogue_msg.translation.y = drogue_t_vec[1]
        drogue_msg.translation.z = drogue_t_vec[2]
        self.lf_drogue_pub.publish(drogue_msg)

       
        if plane_is_visiable:
            cv2.putText(img_gray, "lf_plane:[{:.2f},{:.2f},{:.2f}]".format(plane_t_vec[2],plane_t_vec[0],plane_t_vec[1]), (100,100),cv2.FONT_HERSHEY_SIMPLEX, 4, 255, thickness=5)
        else:
            cv2.putText(img_gray, "lf_plane:Invalid", (100,100),cv2.FONT_HERSHEY_SIMPLEX, 4, 255, thickness=5)
        if drogue_is_visiable:
            cv2.putText(img_gray, "lf_drogue:[{:.2f},{:.2f},{:.2f}]".format(drogue_t_vec[2],drogue_t_vec[0],drogue_t_vec[1]), (100,300),cv2.FONT_HERSHEY_SIMPLEX, 4, 255, thickness=5)
        else:
            cv2.putText(img_gray, "lf_drogue:Invalid", (100,300),cv2.FONT_HERSHEY_SIMPLEX, 4, 255, thickness=5)
        
        cv2.imwrite(f"/home/nuc001/bs_gd_single-master/bag/image_lf/{time.time()}.jpg",img_gray)
        msg=self.bridge.cv2_to_imgmsg(img_gray,"mono8")   #bgr image
        self.lf_img_pub.publish(msg)

        # print(self.x_plane,self.x_drogue)


    def solve(self):
        if not os.path.exists("/home/nuc001/bs_gd_single-master/bag/image_lf"):
            os.makedirs("/home/nuc001/bs_gd_single-master/bag/image_lf")
        while not rospy.is_shutdown():
            try:
                raw_image = self.cam.data_stream[0].get_image()
                numpy_image = raw_image.get_numpy_array()
                if numpy_image is None:
                    continue
                
                self.lf_img_cb(numpy_image)

            except Exception as e:
                print(traceback.format_exc())
            
            
       

if __name__ == '__main__':
    rospy.init_node('bs_solver_node', anonymous=True)

    # port, baud =  '/dev/ttyUSB0', 115200
    # uart = serial.Serial(port,baud,parity=serial.PARITY_NONE)
    bs_solver = BS_Solver()

    rospy.spin()